
# ----------------------------------------------------------------------
# load all relevant packages
# ----------------------------------------------------------------------
rm(list=ls())
require("tidyverse")
require("ggplot2")
set.seed(123)

# Define effect size
effect_size <- 0.15

# Number of simulations
num_simulations <- 500

# Number of observations per simulation
N <- 6000

# Vectors to store p-values
ATE_pvalues <- numeric(num_simulations)
CATE_pvalues <- numeric(num_simulations)
HTE_pvalues <- numeric(num_simulations)


# Perform simulations
for (i in 1:num_simulations) {
  # Generate random data for factors A (shaming), B (target-ally), and C (shamer-ally)
  A <- sample(0:1, N, replace = TRUE)
  B <- sample(0:1, N, replace = TRUE)
  C <- ifelse(A==0, NA, sample(0:1, N, replace=TRUE))
  
  # Generate random data for factors D (target-ally/nonally) and E (shamer-ally/nonally)
  D <- ifelse(B == 1, sample(c("Australia", "Canada", "France", "Japan", "Israel"), N, replace = TRUE), 
              sample(c("China", "Iran", "North Korea", "Russia", "Venezuela"), N, replace = TRUE))
  
  E <- ifelse(is.na(C), NA, 
              ifelse(C == 1, 
                     ave(D, D, FUN = function(x) sample(setdiff(c("Australia", "Canada", "France", "Japan", "Israel"), x), 1)), 
                     ave(D, D, FUN = function(x) sample(setdiff(c("China", "Iran", "North Korea", "Russia", "Venezuela"), x), 1))))
  
  
  # Simulate outcome variable Y based on effect of shaming (H1a, H2a, H3a/H3b)
  ATE <- effect_size * A + rnorm(N)
  
  # Fit linear model and obtain p-value
  model_ATE <- lm(ATE ~ A)
  ATE_pvalues[i] <- summary(model_ATE)$coefficients['A', 'Pr(>|t|)']
  
  
  # Simulate outcome variable Y based on effect of ally|shaming (H1b,H2b)
  CATE <- ifelse(A == 1, effect_size * C, 0) + rnorm(N)
  
  # Fit linear model and obtain p-value
  model_CATE <- lm(CATE ~ C)
  CATE_pvalues[i] <- summary(model_CATE)$coefficients['C', 'Pr(>|t|)']
  
  
  # Simulate outcome variable Y based on interaction effect shaming*ally (H1c,H2c)
  HTE <- effect_size*A*B + rnorm(N)
  
  # Fit linear model and obtain p-value
  model_HTE <- lm(HTE ~ A * B)
  HTE_pvalues[i] <- summary(model_HTE)$coefficients['A:B', 'Pr(>|t|)']
  
}

# Summary of p-values for effect of shaming
p_level_ATE <- sum(ATE_pvalues < 0.05) / num_simulations
cat("Proportion of significant p-values for effect of shaming (p < 0.05):", p_level_ATE, "\n")

# Summary of p-values for conditional effect of ally|shame
p_level_CATE <- sum(CATE_pvalues < 0.05) / num_simulations
cat("Proportion of significant p-values for effect of ally|shaming=1 (p < 0.05):", p_level_CATE, "\n")

# Summary of p-values for interaction term (ally*shame)
p_level_HTE <- sum(HTE_pvalues < 0.05) / num_simulations
cat("Proportion of significant p-values for ally*shaming (p < 0.05):", p_level_HTE, "\n")

# Create data frames for ggplot
data_ATE <- data.frame(p_values = ATE_pvalues, effect = "Effect of Shaming \n ATE=0.15")
data_CATE <- data.frame(p_values = CATE_pvalues, effect = "Effect of Ally|Shaming \n CATE=0.15")
data_HTE <- data.frame(p_values = HTE_pvalues, effect = "Interaction Effect (Shaming*Ally) \n HTE=0.15")

# Combine data frames
all_data <- rbind(data_ATE, data_CATE, data_HTE)
all_data$effect <- factor(all_data$effect, levels = c("Effect of Shaming \n ATE=0.15", 
                                                      "Effect of Ally|Shaming \n CATE=0.15",
                                                      "Interaction Effect (Shaming*Ally) \n HTE=0.15"))

# Plot histogram of p-values using ggplot
p <- ggplot(all_data, aes(x = p_values)) +
  geom_histogram(binwidth = 0.01, fill = "lightblue", color = "black") +
  geom_vline(xintercept = 0.05, linetype = "dashed", color = "red") +
  labs(title = "Distribution of P-values for Different Effects (500 Simulations, N=6000 per simulation)",
       x = "P-value",
       y = "Number of simulations") +
  theme_minimal()+
  theme(plot.title = element_text(hjust = 0.5)) 

# Get summary statistics for annotations
summary_stats <- data.frame(
  effect = c("Interaction Effect (Shaming*Ally) \n HTE=0.15", "Effect of Shaming \n ATE=0.15", "Effect of Ally|Shaming \n CATE=0.15"),
  p_value = c(p_level_HTE, p_level_ATE, p_level_CATE)
)

# Add annotations
p + geom_text(data = summary_stats, aes(x = 0.5, y = 300, label = paste("Power = ", round(p_value, 2))), color = "black", size = 3, vjust = -0.5) +
  facet_wrap(~ factor(effect, levels = c("Effect of Shaming \n ATE=0.15", 
                                         "Effect of Ally|Shaming \n CATE=0.15",
                                         "Interaction Effect (Shaming*Ally) \n HTE=0.15")))

ggsave("Figures/power_simulation.pdf", width = 8, height = 4)
